import os
import requests
from flask import Flask, request, jsonify
import threading
from datetime import datetime
import argparse
import logging

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

app = Flask(__name__)
lock = threading.Lock()  # 线程锁，确保线程安全

def parse_arguments():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description='AnythingLLM API 代理服务')
    parser.add_argument('--host', default='0.0.0.0', help='Flask 服务监听地址')
    parser.add_argument('--api-key', default='W5R61JQ-G4C4J1Q-M9DYJ98-222Z423', help='AnythingLLM API 密钥')
    parser.add_argument('--api-host', default='10.135.172.131', help='训练机器的 IP 地址或主机名')
    parser.add_argument('--api-port', type=int, default=3001, help='训练机器上 AnythingLLM API 的端口')
    parser.add_argument('--workspace', default='apt_test', help='目标工作区 Slug')  # 新增工作区参数
    parser.add_argument('--debug', action='store_true', help='启用调试模式')
    return parser.parse_args()

args = parse_arguments()

# 根据命令行参数设置 AnythingLLM API 配置（包含工作区路径）
anythingllm_api_key = args.api_key
anythingllm_base_url = f"http://{args.api_host}:{args.api_port}/api/workspace/{args.workspace}"  # 包含工作区的 API 路径

# 检查 AnythingLLM API 是否可用
def check_api_availability():
    try:
        # 健康检查接口需包含工作区路径（若服务端支持，若无则调整为通用路径）
        health_url = f"{anythingllm_base_url}/health"
        logger.info(f"正在检查 AnythingLLM API 连接: {health_url}")
        response = requests.get(health_url, timeout=10)
        
        if response.status_code == 200:
            api_info = response.json()
            logger.info(f"成功连接到 AnythingLLM API - 版本: {api_info.get('version', '未知')}")
            return True
        else:
            logger.error(f"API 状态检查失败，状态码: {response.status_code}, 响应: {response.text}")
            return False
    except requests.exceptions.ConnectionError as e:
        logger.error(f"无法连接到 AnythingLLM API: {str(e)}")
        logger.error(f"请确保训练机器 {args.api_host}:{args.api_port} 可访问且服务正在运行")
        return False
    except Exception as e:
        logger.error(f"检查 API 可用性时发生未知错误: {str(e)}")
        return False

# 调用 AnythingLLM 模型进行推理
def call_anythingllm_model(prompt, context=None, temperature=0.7, max_tokens=500):
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {anythingllm_api_key}"
    }
    
    payload = {
        "question": prompt,
        "temperature": temperature,
        "max_tokens": max_tokens
    }
    
    # 如果提供了上下文，添加到请求中
    if context:
        payload["context"] = context
    
    try:
        # 使用线程锁确保线程安全
        with lock:
            logger.debug(f"发送请求到 {anythingllm_base_url}/chat - 提示长度: {len(prompt)}")
            response = requests.post(
                f"{anythingllm_base_url}/chat",  # 工作区对应的聊天接口
                json=payload,
                headers=headers,
                timeout=300  # 设置较长的超时时间，因为模型推理可能需要较长时间
            )
            
        if response.status_code == 200:
            result = response.json()
            logger.debug(f"API 响应成功 - 来源数量: {len(result.get('sources', []))}")
            return result
        else:
            error_msg = response.text if response.text else f"HTTP 错误 {response.status_code}"
            logger.error(f"API 调用失败: {error_msg}")
            raise Exception(f"API 调用失败 (状态码: {response.status_code}): {error_msg}")
            
    except requests.exceptions.Timeout:
        logger.error("API 请求超时 - 模型推理可能需要更长时间或服务不可用")
        raise
    except requests.exceptions.RequestException as e:
        logger.error(f"网络请求错误: {str(e)}")
        raise
    except Exception as e:
        logger.error(f"调用 AnythingLLM 模型时发生未知错误: {str(e)}")
        raise

# 定义 API 接口 - 单轮对话
@app.route('/api/chat', methods=['POST'])
def chat():
    try:
        data = request.json
        if not data or 'prompt' not in data:
            return jsonify({"error": "缺少必要的 'prompt' 字段"}), 400
            
        prompt = data['prompt']
        context = data.get('context')  # 可选上下文
        temperature = data.get('temperature', 0.7)  # 默认温度
        max_tokens = data.get('max_tokens', 500)  # 默认最大生成长度
        
        # 调用模型
        result = call_anythingllm_model(
            prompt=prompt,
            context=context,
            temperature=temperature,
            max_tokens=max_tokens
        )
        
        return jsonify({
            "prompt": prompt,
            "response": result.get("response", ""),
            "sources": result.get("sources", []),
            "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        }), 200
        
    except Exception as e:
        logger.exception("处理单轮对话请求时出错")
        return jsonify({"error": f"处理请求时出错: {str(e)}"}), 500

# 定义 API 接口 - 多轮对话（维护对话历史）
@app.route('/api/conversation', methods=['POST'])
def conversation():
    try:
        data = request.json
        if not data or 'prompt' not in data:
            return jsonify({"error": "缺少必要的 'prompt' 字段"}), 400
            
        # 从请求中获取对话历史（如果有）
        conversation_id = data.get('conversation_id')
        history = data.get('history', [])
        prompt = data['prompt']
        
        # 将当前问题添加到历史中
        history.append({"role": "user", "content": prompt})
        
        # 构建包含历史的完整提示
        full_prompt = "\n".join([f"{msg['role']}: {msg['content']}" for msg in history])
        
        # 调用模型
        result = call_anythingllm_model(full_prompt)
        
        # 将回答添加到历史中
        response = result.get("response", "")
        history.append({"role": "assistant", "content": response})
        
        return jsonify({
            "conversation_id": conversation_id or datetime.now().strftime("%Y%m%d%H%M%S"),
            "prompt": prompt,
            "response": response,
            "history": history,
            "sources": result.get("sources", []),
            "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        }), 200
        
    except Exception as e:
        logger.exception("处理多轮对话请求时出错")
        return jsonify({"error": f"处理对话请求时出错: {str(e)}"}), 500

if __name__ == "__main__":
    # 检查 API 可用性
    if check_api_availability():
        # 启动 Flask 应用（注意端口默认由 Flask 自动设置为 5000，可通过命令行参数修改）
        logger.info(f"Starting Flask server on {args.host}:5000 (工作区: {args.workspace})")
        app.run(host=args.host, port=5000, debug=args.debug)
    else:
        logger.error("无法连接到 AnythingLLM API，应用将不会启动")